# -*- coding: utf-8 -*-
"""
@Time ： 2024/9/18 18:20
@Auth ： 九问
@File ：consumers.py
@IDE ：PyCharm
@Email : 2750826557@qq.com
"""
import json

from asgiref.sync import sync_to_async
from channels.generic.websocket import AsyncWebsocketConsumer
from rest_framework_simplejwt.tokens import UntypedToken
from rest_framework_simplejwt.authentication import JWTAuthentication
from channels.db import database_sync_to_async

from chat.models import Message
from users.models import UserDetail


@database_sync_to_async
def get_user_from_token(validated_token):
    return JWTAuthentication().get_user(validated_token)


class ChatConsumer(AsyncWebsocketConsumer):
    online_users = {}

    @sync_to_async
    def get_user_detail(self, user):
        return UserDetail.objects.get(user=user)

    async def connect(self):
        token = self.scope['query_string'].decode().split('=')[-1]
        try:
            validated_token = UntypedToken(token)
            user = await get_user_from_token(validated_token)
            self.user = user
            user_detail = await self.get_user_detail(user)
            self.avatar = user_detail.avatar if user_detail else None
            if user is None:
                raise ValueError("Invalid token")
        except Exception as e:
            print(f"Token validation error: {e}")
            await self.close()
            return

        self.username = user.first_name
        self.room_name = self.scope["url_route"]["kwargs"]["room_name"]
        self.room_group_name = f"chat_{self.room_name}"
        self.user_id = user.id
        if self.room_group_name not in ChatConsumer.online_users:
            ChatConsumer.online_users[self.room_group_name] = {}

        ChatConsumer.online_users[self.room_group_name][self.username] = {
            "avatar": self.avatar,
            "id": self.user_id
        }
        # Join room group
        await self.channel_layer.group_add(self.room_group_name, self.channel_name)
        await self.accept()

        # Notify others about the new user
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                "type": "user.join",
                "username": self.username,
                "online_users": self.get_online_users(),
            },
        )

    async def disconnect(self, close_code):
        ChatConsumer.online_users[self.room_group_name].pop(self.username, None)
        await self.channel_layer.group_discard(self.room_group_name, self.channel_name)

        # Notify others about the user leaving
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                "type": "user.leave",
                "username": self.username,
                "online_users": self.get_online_users(),
            },
        )

    def get_online_users(self):
        return [
            {"username": username, "avatar": user_info["avatar"], "id": user_info["id"]}
            for username, user_info in ChatConsumer.online_users[self.room_group_name].items()
        ]

    async def receive(self, text_data):
        text_data_json = json.loads(text_data)
        message = text_data_json["message"]
        # 保存消息到数据库
        await database_sync_to_async(Message.objects.create)(
            sender=self.user,
            text=message,
            room_name=self.room_name
        )
        await self.channel_layer.group_send(
            self.room_group_name, {
                "type": "chat.message",
                "message": message,
                "username": self.username,
                "avatar": self.avatar
            }
        )

    async def user_join(self, event):
        await self.send(text_data=json.dumps({
            "type": "user_join",
            "username": event["username"],
            "online_users": event["online_users"],
        }))

    async def user_leave(self, event):
        await self.send(text_data=json.dumps({
            "type": "user_leave",
            "username": event["username"],
            "online_users": event["online_users"],
        }))

    async def chat_message(self, event):
        message_data = {
            "message": event["message"],
            "avatar": event["avatar"],
            "username": event["username"]
        }
        await self.send(text_data=json.dumps(message_data))

